package informed.sgd;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.Random;

import informed.matrix.Utils;
import informed.samples.Sample;
import informed.samples.VectorValueSample;
import cern.colt.matrix.DoubleMatrix1D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.impl.DenseDoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;

public class LatentTransformation {
	
	private DoubleMatrix2D P;
	private DoubleMatrix2D Q;
	private double lambda;
	private int factors;
	private double learningRate;
	
	private double error;
	private int nSamples;
	private int iter;
	
	public LatentTransformation(int n, int m, int f){
		P = new DenseDoubleMatrix2D(n, f);
		Q = new DenseDoubleMatrix2D(f, m);
		factors = f;
	}
	
	private void init(ArrayList<Sample> samples){
		nSamples = samples.size();
		Random r = new Random(System.currentTimeMillis());
		double scale = 0.1;
		for(int f = 0; f < factors; f++){
			for(int n = 0; n < P.rows(); n++){
				P.setQuick(n, f, scale*r.nextDouble());
			}
			for(int m = 0; m < Q.columns(); m++){
				Q.setQuick(f, m, scale*r.nextDouble());
			}
		}
		learningRate = 0.1;
		lambda = 0.0000;
		System.out.println("Initializing P("+P.rows()+","+P.columns()+") and Q("+Q.rows()+","+Q.columns()+"). Learning rate = " + learningRate);
	}
	
	public void run(ArrayList<Sample> samples, int epochs){
		init(samples);
		for(int i = 0; i < epochs; i++){
			Random rand = new Random(System.currentTimeMillis());
			Collections.shuffle(samples, rand);
			for(Sample s: samples){
			//for(int u = 0; u < 400; u++){
				//Sample s = samples.get(u);
				move(s);
				iter++;
				if(iter%10000 == 0)
					System.out.println(iter);
			}
			error *= 1.0/nSamples;
			System.out.println("Iteration " + iter + " error=" + error);
			error = 0;
		}
	}
	
	private void move(Sample x){
		VectorValueSample s = (VectorValueSample)x;
		double gamma = learningRate;
		//if(s.v.cache == null)
			s.v.cache = Utils.multMatrixAndSparseVector(Q, s.v);
		double e = s.y - Utils.multMatrixRowAndVector(P, s.i, s.v.cache);
		error += e*e;
		DoubleMatrix1D gP = gradientP(s, e, s.v.cache);
		//System.out.println("e:" + e + " Q:" + Q.zSum() + " s.v:" + s.v.normL1() + " cache:" + s.v.cache.zSum() + " P:"+P.zSum() + " gP:" + gP.zSum());
		Iterator<Integer> keys = s.v.getIndices();
		while(keys.hasNext()){
			int k = keys.next();
			DoubleMatrix1D gQ = gradientQ(s, e, k);
			for(int f = 0; f < Q.rows(); f++){
				Q.setQuick(f, k, Q.getQuick(f, k) - gamma*gQ.getQuick(f));
			}
		}
		for(int f = 0; f < P.columns(); f++){
			P.setQuick(s.i, f, P.getQuick(s.i, f) - gamma*gP.getQuick(f));
		}
	}
	
	private DoubleMatrix1D gradientP(VectorValueSample s, double e, DoubleMatrix1D Qx){
		DoubleMatrix1D eSumQx = Utils.multByScalar(Qx, -e);
		DoubleMatrix1D lp = Utils.multMatrixRowByScalar(P, s.i, lambda);
		//System.out.println(" >> gradientP ==> eSumQx:" + eSumQx.zSum() + " lp:" + lp.zSum());
		return Utils.addVectors(eSumQx, lp);
	}
	
	private DoubleMatrix1D gradientQ(VectorValueSample s, double e, int k){
		DoubleMatrix1D exp = Utils.multMatrixRowByScalar(P, s.i, -e*s.v.get(k));
		DoubleMatrix1D lq = Utils.multMatrixColByScalar(Q, k, lambda);
		return Utils.addVectors(exp, lq);
	}
	
	public String printSolution(){
		Algebra algebra = new Algebra();
		String s = "";
		s = "\nMatrix P: ";
		s += P.toString();
		s += "\nMatrix Q: ";
		s += Q.toString();
		s += "\nReconstruction: ";
		s += algebra.mult(P, Q).toString();
		return s;
	}

}
